#import os
import random
import PIL.Image as Image
import numpy as np
#import matplotlib.pyplot as plt
from torchvision import transforms
import torch


def image2label(label_data):
    cm2lbl = np.zeros(256 ** 3)
    for i ,cm in enumerate(COLORMAP):
        cm2lbl[(cm[0]*256 + cm[1]) * 256 + cm[2]]= i

    data = np.array(label_data,dtype='int32')
    idx = (data[:,:,0] * 256 + data[:,:,1]) * 256 + data[:,:,2]

    return np.array(cm2lbl[idx],dtype='int64')





def RandomCrop(data_name,label_name,crop_size):

    '''

    :param img_data: 为image对象名称
    :param img_label: 为image_open对象名称
    :param crop_size: 期望的裁剪尺寸
    :return: data和label对应的裁剪之后的图像
    '''

    img_data = Image.open(data_name)
    img_label= Image.open(label_name).convert('RGB')




    h,w = img_data.size
    th,tw = crop_size
    if(h == th and w == tw):
        return img_data,img_label

    print(h,th,w,tw)


    i = random.randint(0,h - th)

    j = random.randint(0,w - tw)
    crop = (i,j,i+th,j+tw)

    img_data = img_data.crop(crop)
    img_label = img_label.crop(crop)


    return img_data,img_label


# data_name = os.path.join('..','../','data/VOC2012/JPEGImages','2007_001834.jpg')
# label_name= os.path.join('..','../','data/VOC2012/SegmentationClass','2007_001834.png')
#
# img1,img2 = RandomCrop(data_name,label_name,(480,320))
#


def img_transforms(data_name,label_name,crop_size):
    img_data,label_data = RandomCrop(data_name,label_name,crop_size)

    img_tfs = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406],
                             std=[0.224,0.224,0.225])
    ])

    img_data = img_tfs(img_data)
    label_data = image2label(label_data)
    label_data = torch.from_numpy(label_data)#此时label_data已经是tensor

    return img_data,label_data






COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
            [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
            [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128],
            [64, 128, 128], [192, 128, 128], [0, 64, 0], [128, 64, 0],
            [0, 192, 0], [128, 192, 0], [0, 64, 128]]

CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
           'dog', 'horse', 'motorbike', 'person', 'potted plant',
           'sheep', 'sofa', 'train', 'tv/monitor']
